/**
 * Copyright 2019-2020 Huawei Technologies Co., Ltd
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#ifndef FUSION_ENGINE_OPTIMIZER_OPS_KERNEL_STORE_FE_OPS_KERNEL_INFO_STORE_H_
#define FUSION_ENGINE_OPTIMIZER_OPS_KERNEL_STORE_FE_OPS_KERNEL_INFO_STORE_H_
#include <map>
#include <memory>
#include <nlohmann/json.hpp>
#include <string>
#include <unordered_map>
#include <vector>
#include "adapter/common/op_store_adapter_manager.h"
#include "common/math_util.h"
#include "common/op_info_common.h"
#include "common/opskernel/ops_kernel_info_store.h"
#include "common/opskernel/ops_kernel_info_types.h"
#include "common/util/json_util.h"
#include "common/util/op_info_util.h"
#include "graph/node.h"
#include "graph/op_desc.h"
#include "graph/types.h"
#include "graph/utils/attr_utils.h"
#include "graph/utils/node_utils.h"
#include "ops_store/op_kernel_info.h"
#include "ops_kernel_store/sub_ops_store.h"
#include "ops_store/sub_op_info_store.h"

namespace fe {
using SubOpsStorePtr = std::shared_ptr<SubOpsStore>;
using OpsKernelInfoStorePtr = std::shared_ptr<ge::OpsKernelInfoStore>;

struct CheckSuportResult {
  bool result;
  int64_t fe_impl_type;
  int64_t ge_impl_type;
};
class FEOpsKernelInfoStore : public ge::OpsKernelInfoStore {
 public:
  explicit FEOpsKernelInfoStore(OpStoreAdapterManagerPtr op_store_adapter_manager_ptr,
                                std::string engine_name = fe::AI_CORE_NAME);

  ~FEOpsKernelInfoStore() override;

  FEOpsKernelInfoStore(const FEOpsKernelInfoStore &) = delete;

  FEOpsKernelInfoStore &operator=(const FEOpsKernelInfoStore &) = delete;

  /*
   * @ingroup fe
   * @brief : Initialize the FEOpsKernelInfoStore, load the op info from the
   * specified .json file
   * @param[in] options: the reserved param
   * @return : SUCCESS/FAILED
   */
  Status Initialize(const map<string, string> &options) override;

  /*
   * @ingroup fe
   * @brief : finalize the FEOpsKernelInfoStore, clear the op info;
   * @param[in] None
   * @return : SUCCESS/FAILED
   */
  Status Finalize() override;

  Status CreateSession(const std::map<std::string, std::string> &session_options) override;

  Status DestroySession(const std::map<std::string, std::string> &session_options) override;

  /*
   * @ingroup fe
   * @brief : get the specified op info from the FEOpsKernelInfoStore;
   * @param[in] infos ： a struct containing the wanted fields of op info;
   * @return: SUCCESS/FAILED
   */
  void GetAllOpsKernelInfo(std::map<std::string, ge::OpInfo> &infos) const override;

  Status GetHighPrioOpKernelInfoPtr(const std::string &op_type, OpKernelInfoPtr &sub_store_op_kernel_info_ptr);
  /*
   * @brief: Query all the FEOpsKernelInfoStores, get the highest priority
   * implement type of the Op;
   * @param[in] OpDesc: op
   * @param[in|out] OpImplType: the implement type of this Op;
   * @return: SUCCESS/FAILED
   */
  Status QueryHighPrioOpImplType(const ge::OpDescPtr &op_desc_ptr, OpImplType &impl_type);

  Status GetAllSubOpsStore(std::map<std::string, SubOpsStorePtr> &all_sub_store_ptr) const;

  bool CheckSupported(const ge::OpDescPtr &op_desc_ptr, std::string &un_supported_reason) const override;

  bool CheckAccuracySupported(const ge::OpDescPtr &op_desc_ptr, std::string &un_supported_reason,
                              bool real_query = false) const override;

  bool CheckSupportedBase(const ge::OpDescPtr &op_desc_ptr, std::string &un_supported_reason, CheckSupportMode mode,
                          bool real_query = false) const;

  std::string GetFEOpsKernelInfoStoreName() const;

  Status CompileOp(vector<ge::NodePtr> &node_vec) override;

  Status CompileOpRun(vector<ge::NodePtr> &node_vec) override;

  bool CheckSupportedByOpsStore(const ge::OpDescPtr &op_desc_ptr, UnSupportedReason &reason,
                                OpImplType &impl_type, CheckSupportMode mode) const;

  Status GetNotSupportedReasonByAttr(uint64_t &reason, std::ostringstream &reason_oss) const;

  Status SetDynamicCustomOpStoreInfo(ge::ComputeGraph &graph);

  /* If the auto mix precision mode is on, we need to  */
  Status UpdateTensorByMixPrecisionMode(ge::OpDesc &op_desc, OpKernelInfoPtr &op_kernel_info_ptr);

  /* check support for trans-nodes by caching result. */
  bool CheckAccuracySupportByCache(const ge::OpDescPtr &op_desc_ptr);

  /* Store check support result for trans-nodes only. */
  Status StoreCheckSuportResultForTransNodes(const ge::OpDescPtr &op_desc_ptr, bool result);

  void SetCheckSupportedStaticFlag(bool stat_flag);

 private:
  bool init_flag_;

  bool custom_or_builtin_exist_flag_;

  std::map<std::string, SubOpsStorePtr> map_all_sub_store_info_;

  std::unordered_map<uint64_t, CheckSuportResult> checksupport_cache_;

  std::string op_kernel_store_type_;

  std::string engine_name_;

  OpStoreAdapterManagerPtr op_store_adapter_manager_ptr_;

  bool check_support_static_flag_;

  Status CompileOpGetTvmJsonInfo(ScopeNodeIdMap &fusion_nodes_map, map<int64_t, string> &scope_json_map);

  Status CompileOpGetImplType(ge::OpDescPtr op_desc_ptr, OpImplType &impl_type);

  Status PrePareCompileParameter(ge::Node &node, string &op_type, OpImplType &impl_type,
                                 std::unordered_map<OpStoreAdapterPtr, vector<PreCompileNodePara>> &node_map);

  bool IsExistInTBECustom(const ge::NodePtr &node_ptr);

  Status GetDefFeOpsStoreInfo(FEOpsStoreInfo &fe_ops_store);

  Status GetOpImplyRealPath(std::string op_imply_relative_path, const OpImplType &op_impl_type,
                            std::string &op_store_real_path, std::string &op_imply_real_path,
                            const ge::NodePtr &node_ptr);

  Status UpdateOpImplyPath(const ge::NodePtr &node_ptr, std::string &op_store_real_path, const OpImplType &op_impl_type,
                           SubOpInfoStorePtr &sub_custom_ops_kernel_ptr);

  Status GetDynamicCustomOpStoreInfoByNode(const ge::NodePtr &node_ptr, vector<std::string> &json_files,
                                           SubOpInfoStorePtr &sub_dyna_custom_ops_store_ptr);

  Status SetDynaCustomOpStoreToAllStore(FEOpsStoreInfo &fe_ops_store,
                                        SubOpInfoStorePtr &sub_dyna_custom_ops_kernel_ptr);

  Status GetAllAtomicCleanNode(ge::NodePtr &node_ptr, vector<ge::NodePtr> &atomic_node_vec);

  Status CompileAndSetKernelNameForAtomicClean(const vector<ge::NodePtr> &node_vec,
                                               vector<ge::NodePtr> &atomic_clean_nodes);

  Status CompileAtomicClean(vector<ge::NodePtr> &node_vec);

  std::vector<uint32_t> CompileGetAtomicOutput(ge::OpDescPtr &op_desc_ptr);

  Status SetAtomicOpAttr(ge::OpDescPtr &op_desc, bool &atomic_node_flag);

  Status CompileSetAtomicCleanWorkSpace(ge::OpDescPtr &op_desc_ptr, vector<int64_t> &work_space,
                                        vector<int64_t> &work_space_bytes);

  bool OpCheckSupport(const ge::OpDescPtr &op_desc_ptr, FEOpsStoreInfo &ops_store,
                      OpStoreAdapterPtr &op_store_adapter_ptr, OpKernelInfoPtr &op_kernel_info_ptr,
                      UnSupportedReason &reason) const;

  bool CheckCustomOp(const ge::OpDescPtr &op_desc_ptr, FEOpsStoreInfo &ops_store, bool &is_custom_op) const;

  void SetGeOp(const ge::OpDescPtr &op_desc_ptr, const std::string &op_filename) const;

  Status PreCompileAndCompile(std::unordered_map<OpStoreAdapterPtr, vector<PreCompileNodePara>> &node_map,
                              const ge::NodePtr &node, ScopeNodeIdMap &fusion_node_map);

  Status CompileSingleOp(ge::NodePtr &node_ptr);

  Status CompileMultipleOp(vector<ge::NodePtr> &node_vec);

  bool IsNeededCompile(ge::OpDescPtr &op_desc_ptr);

  OpKernelInfoPtr GetOpKernelInfoFromOpsStore(const std::string &ops_store_name, const std::string &op_type) const;

  Status GetTbeAdapter(const ge::OpDescPtr &op_desc_ptr, const FEOpsStoreInfo &ops_store,
                       std::ostringstream &reason_oss, OpStoreAdapterPtr &op_store_adapter,
                       SubOpsStorePtr &sub_ops_store_iter) const;

  Status GetOpKernel(const ge::OpDescPtr &op_desc_ptr, const FEOpsStoreInfo &ops_store, OpKernelInfoPtr &op_kernel_ptr,
      std::ostringstream &reason_oss, uint64_t &not_support_reason_id) const;
};
}  // namespace fe

#endif  // FUSION_ENGINE_OPTIMIZER_OPS_KERNEL_STORE_FE_OPS_KERNEL_INFO_STORE_H_
